import os
import logging
from typing import Any, Dict, List, Optional
from langchain.schema import BaseMessage, HumanMessage, AIMessage, SystemMessage
from langchain.llms.base import LLM
from zhipuai import ZhipuAI

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class ZhipuAILLM(LLM):
    """智谱AI LLM 包装器 - 兼容Pydantic版本"""
    
    model_name: str = "glm-4-flash"
    temperature: float = 0.7
    max_tokens: int = 2048
    api_key: str = ""
    
    class Config:
        """Pydantic配置"""
        arbitrary_types_allowed = True
        extra = "ignore"
    
    def __init__(self, **data: Any):
        super().__init__(**data)
        # 不在这里初始化客户端，延迟到第一次调用时
    
    @property
    def _llm_type(self) -> str:
        return "zhipuai"
    
    def _call(self, prompt: str, stop: Optional[List[str]] = None, **kwargs: Any) -> str:
        """调用智谱AI API"""
        # 获取API密钥
        api_key = self.api_key or os.getenv("ZHIPU_API_KEY")
        if not api_key:
            error_msg = "ZHIPU_API_KEY环境变量未设置"
            logger.error(error_msg)
            return error_msg
        
        try:
            # 每次调用都创建新的客户端，避免序列化问题
            client = ZhipuAI(api_key=api_key)
            logger.info(f"调用智谱AI API，模型: {self.model_name}, 输入长度: {len(prompt)}")
            
            response = client.chat.completions.create(
                model=self.model_name,
                messages=[{"role": "user", "content": prompt}],
                temperature=self.temperature,
                max_tokens=self.max_tokens,
            )
            
            result = response.choices[0].message.content
            logger.info(f"智谱AI API调用成功，输出长度: {len(result)}")
            return result
            
        except Exception as e:
            error_msg = f"智谱AI API调用失败: {str(e)}"
            logger.error(error_msg)
            return error_msg
    
    def _convert_message_to_dict(self, message: BaseMessage) -> Dict[str, Any]:
        """转换消息为字典格式"""
        if isinstance(message, HumanMessage):
            return {"role": "user", "content": message.content}
        elif isinstance(message, AIMessage):
            return {"role": "assistant", "content": message.content}
        elif isinstance(message, SystemMessage):
            return {"role": "system", "content": message.content}
        else:
            return {"role": "user", "content": str(message.content)}

class ZhipuAIChatLLM(ZhipuAILLM):
    """支持多轮对话的智谱AI LLM"""
    
    def generate_chat(self, messages: List[BaseMessage], **kwargs: Any) -> str:
        """生成聊天响应"""
        # 获取API密钥
        api_key = self.api_key or os.getenv("ZHIPU_API_KEY")
        if not api_key:
            error_msg = "ZHIPU_API_KEY环境变量未设置"
            logger.error(error_msg)
            return error_msg
        
        try:
            # 每次调用都创建新的客户端
            client = ZhipuAI(api_key=api_key)
            formatted_messages = [self._convert_message_to_dict(msg) for msg in messages]
            logger.info(f"调用智谱AI聊天API，消息数量: {len(formatted_messages)}")
            
            response = client.chat.completions.create(
                model=self.model_name,
                messages=formatted_messages,
                temperature=self.temperature,
                max_tokens=self.max_tokens,
            )
            
            result = response.choices[0].message.content
            logger.info(f"智谱AI聊天API调用成功，输出长度: {len(result)}")
            return result
            
        except Exception as e:
            error_msg = f"调用智谱AI聊天API时出错: {str(e)}"
            logger.error(error_msg)
            return error_msg